package com.lovegu.one;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.List;
import java.util.Scanner;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * <p>
 * 实现端口扫描器
 * </p>
 *
 * @author 老顾
 * @since 2023/3/28
 */
public class PortScanner {

    private static final int TIMEOUT = 2000; // 连接超时时间
    private static final int THREAD_NUM = 10; // 线程池中的线程数量

    public static void main(String[] args) throws InterruptedException {

        Scanner scanner = new Scanner(System.in);
        System.out.print("请输入目标主机IP地址：");
        String host = scanner.nextLine();
        System.out.print("请输入起始端口号：");
        int startPort = scanner.nextInt();
        System.out.print("请输入结束端口号：");
        int endPort = scanner.nextInt();

        int portRange = endPort - startPort + 1;
        int threadNum = Math.min(portRange, THREAD_NUM); // 线程数量不超过端口数目

        ExecutorService executor = Executors.newFixedThreadPool(threadNum);
        // 为了保证线程之间的同步，使用 CountDownLatch 工具类来等待所有线程完成任务
        CountDownLatch latch = new CountDownLatch(portRange);

        List<Integer> openPorts = new ArrayList<>();

        for (int i = 0; i < portRange; i++) {
            int port = startPort + i;
            executor.execute(() -> {
                if (isPortOpen(host, port)) {
                    openPorts.add(port);
                }
                // 每个线程完成任务后，会调用`CountDownLatch`的`countDown()`方法来减少计数器的值
                // 当计数器的值减少到0时，主线程会从等待中返回，并合并所有线程的扫描结果。
                latch.countDown();
            });
        }

        latch.await();

        executor.shutdown();

        System.out.println("扫描结果：");
        if (openPorts.isEmpty()) {
            System.out.println("没有发现开放的端口。");
        } else {
            for (int port : openPorts) {
                System.out.println("端口 " + port + " 开放。");
            }
        }
    }

    private static boolean isPortOpen(String host, int port) {
        try {
            // 为了提高扫描效率，使用了NIO的方式来实现端口扫描
            // 使用 SocketChannel 类来创建非阻塞式的套接字通道
            SocketChannel channel = SocketChannel.open();
            channel.configureBlocking(false);
            channel.connect(new InetSocketAddress(host, port));

            long startTime = System.currentTimeMillis();

            // 在连接目标主机的过程中，使用了超时机制，如果连接超时，则认为该端口不开放
            while (!channel.finishConnect()) {
                if (System.currentTimeMillis() - startTime > TIMEOUT) {
                    channel.close();
                    return false;
                }
            }

            ByteBuffer buffer = ByteBuffer.allocate(1024);
            // 在连接成功后，向目标主机发送一个"Hello"字符串
            buffer.put("Hello".getBytes());
            buffer.flip();
            while (buffer.hasRemaining()) {
                channel.write(buffer);
            }
            buffer.clear();
            channel.read(buffer);
            buffer.flip();
            String response = new String(buffer.array(), 0, buffer.limit());
            channel.close();
            // 如果目标主机返回一个"World"字符串，则认为该端口开放。
            return response.equals("World");
        } catch (IOException e) {
            return false;
        }
    }
}
